#lang racket

(require "sets.rkt")

(provide syntax-identifier? syntax? syntax-pair? syntax-atom?
         syntax-identifier syntax-atom
         syntax-expression syntax-scope
         syntax-add-scope syntax-remove-scopes
         syntax-car syntax-cdr syntax-null?
         print-syntax
         syntax->list
         datum->syntax syntax->datum syntax-cons
         syntax-map)

(define (print-syntax-pair stx port mode)
  (let ([recur (case mode
                 [(#t) write]
                 [(#f) display]
                 [else (lambda (p port) (print p port mode))])])
    (recur (syntax->datum* stx) port)))

(define (print-syntax-identifier stx port mode)
  (let ([recur (case mode
                 [(#t) write]
                 [(#f) display]
                 [else (lambda (p port) (print p port mode))])])
    (recur (syntax-expression stx) port)
    (when mode
      (let ((scope (set->list (syntax-scope stx))))
        (unless (null? scope)
          (write-string "{" port)
          (recur (car scope) port)
          (for-each (lambda (e)
                      (write-string ", " port)
                      (recur e port))
                    (cdr scope))
          (write-string "}" port))))))

(struct syntax-identifier (expression location (scope #:mutable))
  #:methods gen:custom-write
  [(define write-proc print-syntax-identifier)])

(struct syntax-pair ((expression #:mutable) location (scope #:mutable) (pushed? #:mutable))
  #:methods gen:custom-write
  [(define write-proc print-syntax-pair)])

(struct syntax-atom (expression location))

(define (syntax? s)
  (or (syntax-identifier? s) (syntax-atom? s) (syntax-pair? s)))

(define (syntax-null? s) (null? (syntax-expression s)))
(define (syntax-car s) (car (syntax-expression s)))
(define (syntax-cdr s) (cdr (syntax-expression s)))
(define (syntax->list s)
  (let loop ((s s) (m '()))
    (cond ((syntax-null? s) m)
          ((syntax-atom? s) #f)
          ((syntax-pair? s) (loop (syntax-cdr s) (cons (syntax-car s) m))))))

(define (syntax-location s)
  (cond ((syntax-identifier? s) (syntax-identifier-location s))
        ((syntax-atom? s) (syntax-atom-location s))
        ((syntax-pair? s) (syntax-pair-location s))
        (else (error "this object has no scope" s))))

(define (syntax-scope s)
  (cond ((syntax-identifier? s) (syntax-identifier-scope s))
        ((syntax-pair? s) (syntax-pair-scope s))
        (else (error "this object has no scope" s))))

(define (syntax-expression s)
  (cond ((syntax-identifier? s) (syntax-identifier-expression s))
        ((syntax-pair? s)
         (let ((scope (syntax-pair-scope s))
               (expression (syntax-pair-expression s)))
           (if (syntax-pair-pushed? s)
               expression
               (let ((expression* (cons (syntax-add-scope (car expression) scope)
                                        (syntax-add-scope (cdr expression) scope))))
                 (set-syntax-pair-expression! s expression*)
                 (set-syntax-pair-pushed?! s #t)
                 expression*))))
         ((syntax-atom? s) (syntax-atom-expression s))
         (else (error "not a syntax object" s))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (syntax-add-scope s scope)
  (cond ((syntax-identifier? s)
         (syntax-identifier (syntax-expression s)
                            (syntax-location s)
                            (set-union (syntax-scope s) scope)))
        ((syntax-pair? s)
         (syntax-pair (syntax-expression s)
                      (syntax-location s)
                      (set-union (syntax-scope s) scope)
                      #f))
        ((syntax-atom? s) s)
        (else (error "not syntax object" s))))

(define (syntax-remove-scopes s scope)
  (cond ((syntax-identifier? s)
         (syntax-identifier (syntax-expression s)
                            (syntax-location s)
                            (set-subtract (syntax-scope s) scope)))
        ((syntax-pair? s)
         (syntax-pair (syntax-expression s)
                      (syntax-location s)
                      (set-subtract (syntax-scope s) scope)
                      #f))
        ((syntax-atom? s) s)
        (else (error "not syntax object" s))))

(define (syntax-map fn sy)
   (cond
    ((syntax-null? sy) sy)
    ((syntax-atom? sy) (fn sy))
    ((syntax-identifier? sy) (fn sy))
    ((syntax-pair? sy)
     (syntax-pair (cons (fn (syntax-car sy))
                        (syntax-map fn (syntax-cdr sy)))
                  (syntax-location sy) (syntax-scope sy) (syntax-pair-pushed? sy)))
    (else (error "not a syntax object" sy))))

(define (syntax->datum s)
  (cond
   ((syntax-pair? s) (let ((e (syntax-expression s)))
                       (cons (syntax->datum (car e)) (syntax->datum (cdr e)))))
   ((syntax-atom? s) (syntax-expression s))
   ((syntax-identifier? s) (syntax-expression s))
   (else (error "not a syntax object" s))))

(define (syntax->datum* s)
  (cond
   ((syntax-identifier? s) s)
   ((syntax-atom? s) (syntax-expression s))
   ((syntax-pair? s) (let ((e (syntax-expression s)))
                       (cons (syntax->datum* (car e)) (syntax->datum* (cdr e)))))
   (else (error "not a syntax object" s))))

(define (print-syntax s)
  (print (syntax->datum* s)))

(define (datum->syntax s)
  (cond
    ((syntax-pair? s)
     (syntax-pair (cons (datum->syntax (syntax-car s))
                        (datum->syntax (syntax-cdr s)))
                  (syntax-location s) (syntax-scope s) (syntax-pair-pushed? s)))
    ((syntax? s) s)
    ((symbol? s) (syntax-identifier s #f (set)))
    ((pair? s) (syntax-pair (cons (datum->syntax (car s))
                                  (datum->syntax (cdr s)))
                            #f (set) #f))
    (else (syntax-atom s #f))))

(define (syntax-cons x y)
  (syntax-pair (cons (datum->syntax x) (datum->syntax y))
               #f (set) #f))

